import pathlib

import ipywidgets as widgets
import numpy as np
from matplotlib.figure import Figure

%matplotlib inline

volume_0_25 = np.load("out/dataset-id-75492803-faithful-downsample-0-scale-25.npy")
volume_5_25 = np.load("out/dataset-id-75492803-faithful-downsample-5-scale-25.npy")
volume_5_100 = np.load("out/dataset-id-75492803-faithful-downsample-5-scale-100.npy")

volume_0_25.shape, volume_5_25.shape, volume_5_100.shape

fig = Figure((15, 10))
fig.set_tight_layout(True)
ax1, ax2 = fig.subplots(1, 2, sharey=True)

ax1.set_title("downsample=0")
ax1.imshow(volume_0_25[:, 160, :].T, cmap="gray")
ax1.grid(True, linestyle=":", color="gray")

ax2.set_title("downsample=5")
ax2.imshow(volume_0_25[:, 160, :].T, cmap="gray")
ax2.grid(True, linestyle=":", color="gray")

fig

@widgets.interact(
    idx=widgets.IntSlider(value=100, min=0, max=volume_5_25.shape[-1]-1, continuous_update=False),
)
def _(idx):
    fig = Figure((15, 10))
    fig.set_tight_layout(True)
    
    ax1, ax2 = fig.subplots(1, 2, sharey=True)
    ax1.set_title("downsample=0")
    ax1.imshow(volume_0_25[..., idx].T, cmap="gray")
    ax2.set_title("downsample=5")
    ax2.imshow(volume_5_25[..., idx].T, cmap="gray")
    
    return fig

fig = Figure((15, 10))
fig.set_tight_layout(True)
ax1, ax2 = fig.subplots(1, 2)

ax1.set_title("scale=25")
ax1.imshow(volume_5_25[:, 160, :].T, cmap="gray")
ax1.grid(True, linestyle=":", color="gray")

ax2.set_title("scale=100")
ax2.imshow(volume_5_100[:, 40, :].T, cmap="gray")
ax2.grid(True, linestyle=":", color="gray")

fig

@widgets.interact(
    idx=widgets.IntSlider(value=25, min=0, max=volume_5_100.shape[-1]-1, continuous_update=False),
)
def _(idx):
    fig = Figure((15, 10))
    fig.set_tight_layout(True)
    
    ax1, ax2 = fig.subplots(1, 2)
    ax1.set_title("scale=25")
    ax1.imshow(volume_5_25[..., idx * 4].T, cmap="gray")
    ax2.set_title("scale=100")
    ax2.imshow(volume_5_100[..., idx].T, cmap="gray")
    
    return fig